import numpy as np   
import pandas as pd 
import os

class Get_npz_label():
    def __init__(self, args) -> None:
        self.args = args
        if self.args.feature_root is None:
            self.args.feature_root = 'data/cmodty/orgdata/comdty.npz'
        self.data = np.load(self.args.label_root, allow_pickle=True)
        self.datetime_list = pd.to_datetime(self.data['times'])
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_label_reg(self):
        label = self.df_make(self.data['lagRet'])
        label = label.rolling(self.args.label_range).sum()
        if self.args.model_type == 'ml':
            label = label.shift(-self.args.label_range)
        if self.args.divide_vol:
            vol = self.df_make(self.data['vol'])
            label = label/vol
        if self.args.varieties is not None:
            label = label[self.args.varieties]
        label = label[(label.index >= self.beg_time) & (label.index <= self.end_time)]
        codes = label.columns
        times = np.array(label.index)
        label = (label.values)
        label = np.expand_dims(label, axis=2)
        return label, codes, times
    
    def df_make(self, data):
        org_shape = data.shape
        data = data.reshape(-1, org_shape[2])
        data_df = pd.DataFrame(data)
        data_df.index = pd.to_datetime(self.datetime_list)
        data_df.columns = self.data['codes']
        return data_df
    
class Get_label_data():
    def __init__(self, args) -> None:
        self.args = args
        if args.label_root is None:
            self.args.label_root = r'data\two_contract'
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_label_reg(self, label_name = 'lagret.pkl.gzip', vol_name='vol.pkl.gzip'):  
        label = pd.read_pickle(os.path.join(self.args.label_root,label_name)).astype(float)
        label = label.rolling(self.args.label_range, min_periods=1).sum().shift(-self.args.label_range)
        if self.args.divide_vol:
            # vol = pd.read_pickle(os.path.join(self.args.label_root,vol_name))
            vol = pd.read_pickle(os.path.join(self.args.label_root,label_name)).astype(float)
            vol = vol.rolling(100, min_periods=1).std()
            label = label/vol
        if self.args.varieties is not None:
            label = label[self.args.varieties]
        label = label[(label.index >= self.beg_time) & (label.index <= self.end_time)]
        codes = label.columns
        times = np.array(label.index)
        label = (label.values)
        label = np.expand_dims(label, axis=2)
        return label, codes, times  


class Get_class_label():
    def __init__(self, args) -> None:
        self.args = args
        if args.label_root is None:
            self.args.label_root = r'data\two_contract'
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_label_reg(self, label_name = 'lagret.pkl.gzip', vol_name='vol.pkl.gzip'):  
        label = pd.read_pickle(os.path.join(self.args.label_root,label_name)).astype(float)
        label = label.rolling(self.args.label_range,min_periods=1).sum().shift(-self.args.label_range)
        label = (label>self.args.label_ret).astype(float)
        if self.args.varieties is not None:
            label = label[self.args.varieties]
        label = label[(label.index >= self.beg_time) & (label.index <= self.end_time)]
        codes = label.columns
        times = np.array(label.index)
        label = (label.values)
        print('获取特征 label 完成',label.shape, np.isnan(label.astype(float)).sum())
        label = np.expand_dims(label, axis=2)
        return label, codes, times 

class Get_rank_label():
    def __init__(self, args) -> None:
        self.args = args
        if args.label_root is None:
            self.args.label_root = r'data\two_contract'
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_label_reg(self, label_name = 'lagret.pkl.gzip', vol_name='vol.pkl.gzip'):  
        label = pd.read_pickle(os.path.join(self.args.label_root,label_name)).astype(float)
        label = label.rolling(self.args.label_range, min_periods=1).sum().shift(-self.args.label_range)
        if self.args.divide_vol:
            # vol = pd.read_pickle(os.path.join(self.args.label_root,vol_name))
            vol = pd.read_pickle(os.path.join(self.args.label_root,label_name)).astype(float)
            vol = vol.rolling(100, min_periods=1).std()
            label = label/vol
        if self.args.varieties is not None:
            label = label[self.args.varieties]
        label = label[(label.index >= self.beg_time) & (label.index <= self.end_time)]
        label = label.rank(pct=True, axis=1)
        codes = label.columns
        times = np.array(label.index)
        label = (label.values)
        label = np.expand_dims(label, axis=2)
        return label, codes, times  
    